'''筛选一些笨鸟的任务到milvus'''
'''！！！！！！！！注意这个不能上传到仓库！！！！！'''

import datetime

import loguru
import numpy as np
import spacy
from loguru import logger

logger.info('spacy model启动加载')
time1 = datetime.datetime.now()
nlp = spacy.load('zh_core_web_md')
time2 = datetime.datetime.now()
logger.info(f'spacy model加载完成，耗时{time2 - time1}秒')

from tests.pool_conn_utils import pool_util
from tqdm import tqdm

from_conn_list = ['internal.i2soft.net', 'dev_timecardM', 'fBjscGV0bvRGCYW2iUQ8', 'timecarddb', 11627, '']
db_from = pool_util(from_conn_list)

import time

import numpy as np
from pymilvus import (
    connections,
    utility,
    FieldSchema, CollectionSchema, DataType,
    Collection,
)

fmt = "\n=== {:30} ===\n"
search_latency_fmt = "search latency = {:.4f}s"
num_entities, dim = 3000, 300


# 近期任务查询
def personal_task_info():
    task_dic = {}
    task_sql = "SELECT * FROM `tbl_task_target` ORDER BY CREATE_TIME DESC limit 15000"
    task_dic_list = db_from.fetch_all(task_sql)
    task_list = []
    for task in tqdm(task_dic_list):
        TARGET_NAME = task['TARGET_NAME']
        if len(str(TARGET_NAME)) >= 6:
            if TARGET_NAME.find('/') >= 0:
                continue
            task_list.append(TARGET_NAME)
    task_list = list(set(task_list))
    import json
    j_list = {i + 1: x for i, x in enumerate(task_list)}
    with open('bn_task_index.json', 'w', encoding='utf-8') as wf:
        json.dump(j_list, wf, ensure_ascii=False, indent=4)
    return task_list


def emb(task_list):
    result = []
    for task in tqdm(task_list):
        dic = {}
        token_s = nlp(task)
        dic['text'] = task
        dic['emb'] = token_s.vector
        result.append(dic)
    return result


def upload(info_dic_list):
    print(fmt.format("start connecting to bn_task"))
    connections.connect("default", host="81.68.223.19", port="19530")
    #
    has = utility.has_collection("bn_task")
    print(f"Does collection bn_task exist in bn_task: {has}")
    #
    fields = [
        FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=False),
        FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=dim)
    ]
    schema = CollectionSchema(fields, "bn_task is the simplest demo to introduce the APIs")
    print(fmt.format("Create collection `bn_task`"))
    bn_task = Collection("bn_task", schema, consistency_level="Strong")
    #
    print(fmt.format("Start inserting entities"))
    entities = [
        # provide the pk field because `auto_id` is set to False
        [i + 1 for i in range(len(info_dic_list))],
        [x['emb'].tolist() for x in info_dic_list]
    ]
    insert_result = bn_task.insert(entities)
    print(f"Number of entities in bn_task: {bn_task.num_entities}")  # check the num_entites
    #
    print(fmt.format("Start Creating index IVF_FLAT"))
    index = {
        "index_type": "IVF_FLAT",
        "metric_type": "L2",
        "params": {"nlist": 300},
    }
    bn_task.create_index("embeddings", index)


if __name__ == '__main__':
    task_list = personal_task_info()
    info_dic_list = emb(task_list)
    upload(info_dic_list)
